#pragma once

#include "socket.h"
#include <algorithm>
#include <bitset>
#include <cstddef>
#include <cstdint>
#include <map>
#include <string>
#include <vector>

namespace cobt {
constexpr uint8_t ID_LENGTH = 160;
struct DHTNode {
  DHTNode() {}
  DHTNode(const std::bitset<160> &id, const SocketAddress &addr,
          bool peer = false)
      : m_id(id), m_addr(addr), m_id_str(ConvertToString(id)), m_peer(peer) {}
  DHTNode(const std::string &id, const SocketAddress &addr, bool peer = false)
      : m_id(ConvertToBitset(id)), m_addr(addr), m_id_str(id), m_peer(peer) {}
  DHTNode(const DHTNode &node) = default;
  DHTNode(DHTNode &&node) = default;
  DHTNode &operator=(const DHTNode &node) = default;
  DHTNode &operator=(DHTNode &&node) = default;

  std::bitset<160> m_id; // 节点的ID (通常是一个160位的哈希值)
  std::string m_id_str;
  SocketAddress m_addr;
  bool m_peer = false;
  std::string m_info_hash;

  std::string compact_info() const {
    if (m_addr.get_family() == AF_INET) {
      std::string compact_node(26, 0);
      memcpy(compact_node.data(), m_id_str.data(), 20);
      memcpy(compact_node.data() + 20, &(m_addr.get_addr4()->sin_addr), 4);
      memcpy(compact_node.data() + 24, &(m_addr.get_addr4()->sin_port), 2);
      return compact_node;
    }
    std::string compact_node(38, 0);
    memcpy(compact_node.data(), m_id_str.data(), 20);
    memcpy(compact_node.data() + 20, &(m_addr.get_addr6()->sin6_addr), 16);
    memcpy(compact_node.data() + 36, &(m_addr.get_addr6()->sin6_port), 2);
    return compact_node;
  }
  static std::bitset<ID_LENGTH> ConvertToBitset(const std::string &buf) {
    std::bitset<ID_LENGTH> bitset;
    for (int i = 0; i < ID_LENGTH / 8; ++i) {
      char byte = buf[i];
      for (int j = 0; j < 8; ++j) {
        bool bit = (byte >> j) & 0x01;
        bitset.set(i * 8 + j, bit);
      }
    }
    return bitset;
  }
  static std::string ConvertToString(const std::bitset<ID_LENGTH> &bitset) {

    constexpr int buf_size = ID_LENGTH / 8;
    std::string buf;
    buf.resize(buf_size);
    for (size_t i = 0; i < buf_size; ++i) {
      char byte = 0;
      for (size_t bit = 0; bit < 8; ++bit) {
        if (bitset[i * 8 + bit]) {
          byte |= (1 << bit);
        }
      }
      buf[i] = byte;
    }
    return buf;
  }
  bool operator==(const DHTNode &other) const { return m_id == other.m_id; }
};

class RoutingTable {
  friend class DHTDigger;

private:
  std::vector<std::vector<DHTNode>> m_buckets;
  DHTNode m_node;
  int m_k;

public:
  RoutingTable(const DHTNode &node = {}) : m_node(node) {
    // 初始化路由表，这里假设160位的ID，因此有160个桶
    m_buckets.resize(160);
  }

  bool insert(const DHTNode &node) {
    int bucket_idx = get_bucket_index(node.m_id);
    auto &bucket = m_buckets[bucket_idx];
    auto it =
        std::find_if(bucket.begin(), bucket.end(),
                     [&node](const DHTNode &n) { return n.m_id == node.m_id; });

    if (it != bucket.end()) {
      // 如果节点已经存在，将其移到列表的末尾（最近使用）
      bucket.erase(it);
      bucket.push_back(node);
      return false;
    } else {
      if (bucket.size() < m_k) {
        bucket.push_back(node); // 插入节点到k桶末尾
      } else {
        // 桶已满
        bucket.front() = node;
        return false;
      }
      return true;
    }
  }

  struct BitsetComparator {
    bool operator()(const std::bitset<ID_LENGTH> &bset1,
                    const std::bitset<ID_LENGTH> &bset2) const {
      for (int i = ID_LENGTH - 1; i >= 0; --i) { // 从高位到低位进行比较
        if (bset1[i] != bset2[i]) {
          return bset1[i] < bset2[i];
        }
      }
      return false; // 如果完全相等，则不小于
    }
  };

  std::vector<DHTNode> closest_nodes(const std::bitset<ID_LENGTH> &targetId) {
    std::map<std::bitset<ID_LENGTH>, DHTNode, BitsetComparator> distances;
    // 找到离目标ID最近的k个节点
    std::vector<DHTNode> res;
    for (auto &bucket : m_buckets) {
      for (auto &node : bucket) {
        distances[get_distance(targetId, node.m_id)] = node;
      }
    }
    int i = 0;
    for (auto &pair : distances) {
      res.push_back(pair.second);
      if (++i >= m_k - 1) {
        break;
      }
    }
    res.push_back(m_node);
    return res;
  }

  std::vector<DHTNode> closest_nodes(const std::string &targetId) {
    if (targetId.size() != ID_LENGTH / 8) {
      return {};
    }
    return closest_nodes(DHTNode::ConvertToBitset(targetId));
  }

private:
  std::bitset<ID_LENGTH> get_distance(const std::bitset<ID_LENGTH> &id1,
                                      const std::bitset<ID_LENGTH> &id2) {
    // 计算XOR距离
    return (id1 ^ id2);
  }

  int get_bucket_index(const std::bitset<ID_LENGTH> &id) const {
    std::bitset<ID_LENGTH> xor_res = m_node.m_id ^ id;
    for (int i = ID_LENGTH - 1; i >= 0; --i) {
      if (xor_res[i]) {
        return ID_LENGTH - 1 - i;
      }
    }
    return 0; // ID相同的情况
  }
};
} // namespace cobt